/**
 * Copyright (c) 2022 Presto Labs Pte. Ltd.
 * Author: Ju
 */

#include "coin2/controller/message_assembler.h"

#include <glog/logging.h>

namespace {
static const uint16_t kLengthBytes = sizeof(uint32_t);

static uint32_t GetBufferLength(const coin2::base::memory::MemSlice &buffer) {
  return *reinterpret_cast<const uint32_t *>(buffer.Begin());
}
}  // namespace
namespace coin2::controller {

MessageAssembler::MessageAssembler(Callback *callback)
    : queueing_bytes_(0), callback_(callback) {}

void MessageAssembler::OnPacketReceived(memory::MemSlice payload) {
  slices_.emplace_back(payload);
  queueing_bytes_ += payload.GetUsedSize();
  CheckCompleteness();
}

bool MessageAssembler::GetPacketLength(uint16_t *length) {
  if (queueing_bytes_ < kLengthBytes) {
    return false;
  }
  if (slices_.front().GetUsedSize() >= kLengthBytes) {
    *length = GetBufferLength(slices_.front());
  } else {
    memory::MemSlice buffer(memory::MemBuf::Create(kLengthBytes), 0);
    uint32_t bytes_needed = kLengthBytes;
    auto it = slices_.begin();
    while (it != slices_.end()) {
      uint32_t bytes = std::min(bytes_needed, it->GetUsedSize());
      memcpy(const_cast<void *>(buffer.End()), it->Begin(), bytes);
      buffer.AdjustEndPointer(bytes);
      if (bytes == bytes_needed) {
        break;
      } else {
        bytes_needed -= bytes;
        ++it;
      }
    }
    *length = GetBufferLength(buffer);
  }
  return true;
}

void MessageAssembler::CheckCompleteness() {
  uint16_t length = 0;
  while (GetPacketLength(&length)) {
    if (length < kLengthBytes && callback_) {
      callback_->OnBroken();
      break;
    }
    if (queueing_bytes_ >= length) {
      auto buffer = TakeBuffer(length);
      if (callback_) {
        
        callback_->OnProcessMessage(std::move(buffer));
      }
    } else {
      LOG(INFO) << "Break";
      break;
    }
  }
}

memory::MemSlice MessageAssembler::TakeBuffer(uint32_t bytes_needed) {
  queueing_bytes_ -= bytes_needed;
  if (bytes_needed <= slices_.front().GetUsedSize()) {
    memory::MemSlice buffer = slices_.front();
    buffer.SetRelativeEndPointer(bytes_needed);
    slices_.front().AdjustBeginPointer(bytes_needed);
    if (slices_.front().GetUsedSize() == 0) {
      slices_.pop_front();
    }
    return buffer;
  }
  memory::MemSlice buffer(memory::MemBuf::Create(bytes_needed), 0);
  auto it = slices_.begin();
  while (it != slices_.end() && bytes_needed > 0) {
    uint32_t bytes = std::min(bytes_needed, it->GetUsedSize());
    memcpy(const_cast<void *>(buffer.End()), it->Begin(), bytes);
    buffer.AdjustEndPointer(bytes);
    it->AdjustBeginPointer(bytes);
    if (it->GetUsedSize() == 0) {
      it = slices_.erase(it);
    }
    bytes_needed -= bytes;
  }
  return buffer;
}

}  // namespace coin2::controller